{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 额外安装一个库  pip install rouge-chinese"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '/media/dengyunfei/6T/data/datasets/supremezxc/nlpcc_data.json'\n",
    "model_dir = '/media/dengyunfei/6T/data/models/huggingface/mengzi-t5-base'\n",
    "save_dir = '/media/dengyunfei/6T/data/logs/abstract_model'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dengyunfei/miniconda3/envs/torch24/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "这里是False\n"
     ]
    }
   ],
   "source": [
    "# 导入相关包\n",
    "import torch\n",
    "from datasets import Dataset,load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,T5ForConditionalGeneration,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating train split: 50000 examples [00:02, 21867.97 examples/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['title', 'content'],\n",
      "    num_rows: 5000\n",
      "})\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# 加载数据集\n",
    "datasets = load_dataset('json', data_files=data_dir,split='train')\n",
    "datasets = datasets.train_test_split(test_size=0.90)\n",
    "datasets = datasets[\"train\"]\n",
    "print(datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['title', 'content'],\n",
       "        num_rows: 4900\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['title', 'content'],\n",
       "        num_rows: 100\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets = datasets.train_test_split(100,seed=42)\n",
    "datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'title': '末代皇帝溥仪弟弟爱新觉罗·溥任今日下午去世,享年97岁;抗战胜利后在北京办小学,系北京市第七、八、九届政协委员。',\n",
       " 'content': '法制晚报讯(记者钱业)记者从溥任家人独家获悉,末代皇帝溥仪的弟弟溥任今天下午三点去世,享年97岁。爱新觉罗·溥任,又名金友之,1918年9月生于北京什刹海北岸摄政王府。1947年他创办北京竞业小学,至1968年退休。曾任北京市第七、八、九届政协委员,北京市文史研究馆馆员。(记者钱业)'}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n"
     ]
    }
   ],
   "source": [
    "# 数据处理\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_func(examples):\n",
    "    contents = [\"摘要生成:\\n\" + e for e in examples[\"content\"]]\n",
    "    # 分开处理\n",
    "    inputs =  tokenizer(contents, max_length=512, truncation=True)\n",
    "    labels =  tokenizer(text_target=examples[\"title\"], max_length=128, truncation=True)\n",
    "    inputs[\"labels\"] =  labels[\"input_ids\"]\n",
    "    return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 4900/4900 [00:01<00:00, 3765.14 examples/s]\n",
      "Map: 100%|██████████| 100/100 [00:00<00:00, 3243.90 examples/s]\n"
     ]
    }
   ],
   "source": [
    "tokenizer_ds =  datasets.map(process_func,batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "摘要生成: 法制晚报讯(记者钱业)记者从溥任家人独家获悉,末代皇帝溥仪的弟弟溥任今天下午三点去世,享年97岁。爱新觉罗·溥任,又名金友之,1918年9月生于北京什刹海北岸摄政王府。1947年他创办北京竞业小学,至1968年退休。曾任北京市第七、八、九届政协委员,北京市文史研究馆馆员。(记者钱业)</s>\n",
      "末代皇帝溥仪弟弟爱新觉罗·溥任今日下午去世,享年97岁;抗战胜利后在北京办小学,系北京市第七、八、九届政协委员。</s>\n"
     ]
    }
   ],
   "source": [
    "# 将数据处理成如下这个样子\n",
    "print(tokenizer.decode(tokenizer_ds[\"train\"][0][\"input_ids\"]))\n",
    "print(tokenizer.decode(tokenizer_ds[\"train\"][0][\"labels\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建模型\n",
    "model =  AutoModelForSeq2SeqLM.from_pretrained(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'rouge_chinese'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[12], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# 创建评估函数\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m----> 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mrouge_chinese\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Rouge\n\u001b[1;32m      4\u001b[0m rouge \u001b[38;5;241m=\u001b[39m Rouge()\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_metric\u001b[39m(eval_preds):\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'rouge_chinese'"
     ]
    }
   ],
   "source": [
    "# 创建评估函数\n",
    "import numpy as np\n",
    "from rouge_chinese import Rouge\n",
    "rouge = Rouge()\n",
    "\n",
    "def compute_metric(eval_preds):\n",
    "    predictions,labels = eval_preds\n",
    "    decoded_preds = tokenizer.batch_decode(predictions,skip_special_tokens=True)\n",
    "    # labels 有-100 因此需要进行填充操作,因为labels有-100 不属于特殊字符，因此需要填充为特殊字符\n",
    "    labels = np.where(labels!=-100,labels,tokenizer.pad_token_id)\n",
    "    decoded_labels = tokenizer.batch_decode(labels,skip_special_tokens=True)\n",
    "    # 基于中文的字来做评估\n",
    "    decoded_preds = [\" \".join(pred.strip()) for pred in decoded_preds]\n",
    "    decoded_labels = [\" \".join(label.strip()) for label in decoded_labels]\n",
    "    # 该方法计算 Rouge，标准的文字评估方案\n",
    "    scores = rouge.get_scores(decoded_preds,decoded_labels,avg=True)\n",
    "    return {\n",
    "        \"rouge-1\":scores[\"rouge-1\"][\"f\"],\n",
    "        \"rouge-2\":scores[\"rouge-2\"][\"f\"],\n",
    "        \"rouge-l\":scores[\"rouge-l\"][\"f\"]\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Seq2SeqTrainingArguments(\n",
    "    output_dir=save_dir,\n",
    "    per_device_train_batch_size=4,\n",
    "    per_device_eval_batch_size=4,\n",
    "    gradient_accumulation_steps=2,\n",
    "    logging_steps = 100,\n",
    "    num_train_epochs=4,\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    metric_for_best_model=\"rouge-l\",\n",
    "    # predict_with_generate 这个参数必须设置为True，否则无法做评估\n",
    "    predict_with_generate = True\n",
    ")\n",
    "args.num_train_epochs              # total number of training epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer =  Seq2SeqTrainer(\n",
    "    args = args,\n",
    "    model =model,\n",
    "    train_dataset = tokenizer_ds[\"train\"],\n",
    "    eval_dataset = tokenizer_ds[\"test\"],\n",
    "    compute_metrics = compute_metric,\n",
    "    tokenizer =tokenizer,\n",
    "    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer)\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开始训练\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\Miniconda\\envs\\geo\\lib\\site-packages\\numpy\\_distributor_init.py:30: UserWarning: loaded more than 1 DLL from .libs:\n",
      "d:\\Miniconda\\envs\\geo\\lib\\site-packages\\numpy\\.libs\\libopenblas.FB5AE2TYXYH2IJRDKGDGQ3XBKLKTF43H.gfortran-win_amd64.dll\n",
      "d:\\Miniconda\\envs\\geo\\lib\\site-packages\\numpy\\.libs\\libopenblas64__v0.3.21-gcc_10_3_0.dll\n",
      "  warnings.warn(\"loaded more than 1 DLL from .libs:\"\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n"
     ]
    }
   ],
   "source": [
    "# 模型推理\n",
    "# 导入相关包\n",
    "import torch\n",
    "import os\n",
    "from datasets import Dataset,load_dataset\n",
    "from transformers import pipeline\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,T5ForConditionalGeneration,DataCollatorForSeq2Seq,Seq2SeqTrainer,Seq2SeqTrainingArguments\n",
    "\n",
    "# data_dir = '/data/datasets/supremezxc/nlpcc_data.json'\n",
    "model_dir = '/data/models/huggingface/mengzi-t5-base'\n",
    "save_dir = '/data/logs/abstract_model'\n",
    "# E:\\data\\logs\\abstract_model\\checkpoint-4900\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(os.path.join(save_dir,'checkpoint-4900')).to('cuda')\n",
    "pipe = pipeline(\"text2text-generation\",model=model,tokenizer=tokenizer,device=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe(\"摘要生成:\\n\"+datasets[\"test\"][-2][\"content\"],max_length=128,do_sample=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = '10月1日晚8点20，WTT中国大满贯男单16强开赛，王楚钦对阵丹麦选手安德斯·林德。这是王楚钦本次大赛的重要一战，他在之前的预选赛中已成功晋级。比赛首局，王楚钦11-13安德斯·林德，比分0:1落后；第二句，王楚钦扳回一局；第三局，王楚钦6-11安德斯·林德，比分1-2落后；第四局林德上场后连续进攻取得领先优势。随后被王楚钦反超比分。然而关键时刻王楚钦自己也出现打丢的情况，叫出暂停也未能调整好心态，最终以7-11无缘晋级男单十六强。'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'generated_text': '2015-07-2713:20新浪体育显示图片王楚钦VS丹麦选手安德斯·林德,王楚钦&林德(右)比赛首局,王楚钦和王楚钦在预选赛中打进两局,王楚钦以6-11领先。第四局,王楚钦11-13安德斯、林德,王楚钦扳回一局。王楚钦扳回一局,王楚钦扳回一局,王楚钦扳回一局。(右)王楚钦换下。王楚钦抢下一局'}]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipe(\"摘要生成:\\n\"+text,max_length=128,do_sample=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch24",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
